from ultralytics import YOLO
import cv2
import time
import numpy as np

CLASS_NAMES = [
    "left",
    "right",
    "back",
    "no_right",
    "roundabout",
    "rateLimit",
    "no_rateLimit",
    "redlight",
    "greenlight",
]

yolo = YOLO("best.pt")

cap = cv2.VideoCapture(0)

def print_max_conf_class(confs, clss):
    if len(confs) == 0:
        print("未检测到目标")
        return
    max_idx = confs.argmax()
    print(f"当前帧置信度最高类别: {CLASS_NAMES[clss[max_idx]]}, 置信度: {confs[max_idx]:.2f}")

def generate_command(confs, clss):
    """
    生成识别指令。
    """
    if len(confs) == 0: 
        return None
    main_cls = clss[np.argmax(confs)]
    return CLASS_NAMES[main_cls]

def recognize():
    # 类别名

    while cap.isOpened():
        start_time = time.time()
        ret, frame = cap.read()
        if not ret:
            break

        results = yolo.predict(frame, conf=0.5, verbose=False)
        annotated_frame = frame.copy()
        boxes = results[0].boxes.xyxy.cpu().numpy()
        confs = results[0].boxes.conf.cpu().numpy()
        clss = results[0].boxes.cls.cpu().numpy().astype(int)

        # 打印当前帧置信度最高的类别
        # print_max_conf_class(confs, clss)

        # 在图像上绘制检测结果
        for box, conf, cls in zip(boxes, confs, clss):
            x1, y1, x2, y2 = map(int, box)
            label = f"{CLASS_NAMES[cls]}:{conf:.2f}"
            cv2.rectangle(annotated_frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
            cv2.putText(annotated_frame, label, (x1, y1-10),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2)

        # 计算帧率
        fps = 1.0 / (time.time() - start_time)
        cv2.putText(annotated_frame, f"FPS: {fps:.2f}", (10, 30),
                    cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)

        cv2.imshow("YOLOv8 Detect", annotated_frame)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    cap.release()
    cv2.destroyAllWindows()